package service

import (
	"context"
	"crypto/rand"
	"fmt"
	"os"

	"github.com/tetratelabs/wazero"
	"github.com/tetratelabs/wazero/api"
	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
	"github.com/tokentransfer/node/core"
	"github.com/tokentransfer/node/storage"
	"github.com/tokentransfer/node/util"
	"github.com/tokentransfer/node/vm/env"
)

type WasmModule struct {
	s core.Storage

	runtimeConfig wazero.RuntimeConfig
	moduleConfig  wazero.ModuleConfig
	runtime       wazero.Runtime

	mod api.Module

	mallocFn  api.Function
	freeFn    api.Function
	startFn   api.Function
	processFn api.Function
	stopFn    api.Function
}

func (wm *WasmModule) Close(cost *api.Cost) error {
	ctx := context.WithValue(context.Background(), "cost", cost)
	return wm.runtime.Close(ctx)
}

func (wm *WasmModule) Load(cost *api.Cost, wasmCode []byte) (api.Module, error) {
	ctx := context.WithValue(context.Background(), "cost", cost)
	runtimeConfig := wazero.NewRuntimeConfigInterpreter()
	runtime := wazero.NewRuntimeWithConfig(ctx, runtimeConfig)

	wasi_snapshot_preview1.MustInstantiate(ctx, runtime)
	moduleConfig := wazero.NewModuleConfig().WithStartFunctions().
		WithStdout(os.Stdout).WithStderr(os.Stderr).
		WithRandSource(rand.Reader)

	if wm.s != nil {
		nfs, err := storage.NewFs(wm.s, "/temp")
		if err != nil {
			return nil, err
		}
		moduleConfig = moduleConfig.WithNewFS(nfs)
	}

	//load dependency modules
	err := env.LoadModule(ctx, runtime)
	if err != nil {
		return nil, err
	}
	// err = wasi.LoadModule(ctx, runtime)
	// if err != nil {
	// 	return nil, err
	// }
	// err = wbg.LoadModule(ctx, runtime)
	// if err != nil {
	// 	return nil, err
	// }
	wm.runtimeConfig = runtimeConfig
	wm.moduleConfig = moduleConfig
	wm.runtime = runtime

	mod, err := runtime.InstantiateWithConfig(ctx, wasmCode, moduleConfig)
	if err != nil {
		return nil, err
	}
	return mod, nil
}

func (wm *WasmModule) Start(cost *api.Cost) error {
	ctx := context.WithValue(context.Background(), "cost", cost)
	_, err := wm.startFn.Call(ctx)
	if err != nil {
		return err
	}
	return nil
}

func (wm *WasmModule) Stop(cost *api.Cost) error {
	ctx := context.WithValue(context.Background(), "cost", cost)
	_, err := wm.stopFn.Call(ctx)
	if err != nil {
		return err
	}
	return nil
}

func (wm *WasmModule) malloc(ctx context.Context, size uint64) (uint64, error) {
	dataPtr, err := wm.mallocFn.Call(ctx, size)
	if err != nil {
		return 0, err
	}
	if len(dataPtr) != 1 {
		return 0, util.ErrorOfInvalid("call malloc", "method")
	}
	return dataPtr[0], nil
}

func (wm *WasmModule) free(ctx context.Context, ptr uint64, size uint32) error {
	_, err := wm.freeFn.Call(ctx, ptr, uint64(size))
	if err != nil {
		fmt.Println(">>> free", err)
		return err
	}
	return nil
}

func (wm *WasmModule) Process(cost *api.Cost, input string) (string, error) {
	ctx := context.WithValue(context.Background(), "cost", cost)
	data := []byte(input)
	dataSize := uint64(len(data))
	dataPtr, err := wm.malloc(ctx, dataSize)
	if err != nil {
		return "", err
	}
	ok := wm.mod.Memory().Write(uint32(dataPtr), data)
	if !ok {
		return "", util.ErrorOfInvalid("write memory", fmt.Sprintf("%d, %d", dataPtr, dataSize))
	}
	defer wm.free(ctx, dataPtr, uint32(dataSize))

	results, err := wm.processFn.Call(ctx, api.EncodeI64(int64(dataPtr<<32|dataSize)))
	if err != nil {
		return "", err
	}
	if len(results) != 1 {
		return "", util.ErrorOfInvalid("call process", "function")
	}
	retPtr, retSize := uint32(results[0]>>32), uint32(results[0])
	retData, ok := wm.mod.Memory().Read(retPtr, retSize)
	if !ok {
		return "", util.ErrorOfInvalid("read memory", fmt.Sprintf("%d, %d", retPtr, retSize))
	}
	// defer wm.free(ctx, uint64(retPtr), retSize)

	return string(retData), nil
}
